#pragma once

#include <optional>

#include "../utils/pack.h"
#include "../utils/misc.h"

namespace mdtx::online::sampler
{
    namespace detail
    {
        template <typename Time, typename T, size_t N>
        struct sampler_bar_impl
        {
            using time_part = std::tuple<Time>;
            using data_part = ntuple<T, N>;
            using type = typename mdtx::online::detail::tuple_join<time_part, data_part>::type;
        };

        template <typename Time, typename...TUPLE>
        inline auto cat_time_tups(Time t, TUPLE &&...tups)
        {
            return std::tuple_cat(std::make_tuple(t), std::forward<TUPLE>(tups)...);
        }
    } // namespace detail

    /**
     * @brief Sampler bar type
     *
     * @tparam Time time type
     * @tparam T scala type
     * @tparam N number of scala fields (sans time)
     */
    template <typename Time, typename T, size_t N>
    using sampler_bar = typename detail::sampler_bar_impl<Time, T, N>::type;

    /**
     * @brief Sampler base
     *
     * @tparam Time time type
     * @tparam T scala type
     * @tparam Nin number of input scala fields (sans time)
     * @tparam Nout number of output scala fields (sans time)
     */
    template <typename Time, typename T, size_t Nin, size_t Nout>
    class sampler_base
    {
    public:
        // time type
        using time_type = Time;
        // input/output scala type
        using scala_type = T;
        // bar type
        using bar_type = sampler_bar<Time, T, Nout>;
        // input tuple
        using tuple_in = sampler_bar<Time, T, Nin>;
        // output optional tuple
        using tuple_out = std::optional<bar_type>;
        // number of input, including time
        constexpr static size_t num_in = Nin + 1;
        // number of input, excluding time
        constexpr static size_t num_vin = Nin;

        // number of output, including time
        constexpr static size_t num_out = Nout + 1;
        // number of output, excluding time
        constexpr static size_t num_vout = Nout;

    protected:
        constexpr static make_index_range<0, num_in> inidx{};
        constexpr static make_index_range<num_in, num_in + num_out> outidx{};

        template <typename TUPLE, size_t... Ins, size_t... Outs>
        inline bool call_update_iter(TUPLE &tup, std::integer_sequence<size_t, Ins...>, std::integer_sequence<size_t, Outs...>)
        {
            tuple_out opt_tup = update(tuple_in(*std::get<Ins>(tup)++...));
            if (opt_tup.has_value())
            {
                std::tie(*std::get<Outs>(tup)++...) = *opt_tup;
                return true;
            }
            else
            {
                return false;
            }
        }

    public:
        virtual tuple_out update(tuple_in const &) = 0;
        virtual tuple_out flush() = 0;

        template <typename... Ts>
        tuple_out update(Ts... in)
        {
            static_assert(sizeof...(Ts) == num_in, "Incorrect number of input.");
            return update(tuple_in(in...));
        }

        template <typename... Iter>
        size_t update_n(size_t n, Iter... inout)
        {
            static_assert(sizeof...(Iter) == (num_in + num_out), "Incorrect number of input/output.");
            auto tup = std::make_tuple(inout...);
            size_t nsampled{};
            for (size_t i = 0; i < n; ++i)
            {
                if (call_update_iter(tup, inidx, outidx))
                {
                    nsampled += 1;
                }
            }
            return nsampled;
        }

        template <typename... Ts>
        tuple_out operator()(Ts... in)
        {
            static_assert(sizeof...(Ts) == num_in, "Incorrect number of input.");
            return update(tuple_in(in...));
        }

        tuple_out operator()(tuple_in const &in)
        {
            return update(in);
        }

        virtual ~sampler_base() = default;
    };
}
